#pragma once

#include <Eigen/Core>
#include "stair_ctrl/math.hpp"
#include "stair_ctrl/kinematics.hpp"
#include <array>


template <typename T>
class robotInterface
{
public:
    Eigen::Matrix<T, 12, 1> joint_pos, joint_vel, joint_vel_filtered, joint_pos_des, joint_kp, joint_kd;
    Eigen::Matrix<T, 3, 1> body_ang_pos, body_ang_vel, body_lin_acc;
    Eigen::Matrix<T, 4, 1> body_quat; // w x y z
    Eigen::Matrix<int16_t, 4, 1> foot_force;
    std::array<bool, 4> contact_status = {true, true, true, true};
    Eigen::Matrix<T, 3, 1> body_lin_pos, body_lin_vel;
    T joint_vel_filter_freq, joint_vel_filter_alpha;

    robotInterface()
    {
        joint_pos.setZero();
        joint_vel.setZero();
        joint_vel_filtered.setZero();
        joint_pos_des.setZero();
        joint_kp.setZero();
        joint_kd.setZero();
        body_ang_pos.setZero();
        body_ang_vel.setZero();
        body_quat.setZero();
        foot_force.setZero();
        body_lin_acc.setZero();

        body_lin_pos.setZero();
        body_lin_vel.setZero();

        joint_vel_filter_freq = 50.0;
        joint_vel_filter_alpha = 20.0 * M_PI * 0.01 * joint_vel_filter_freq / (1.0 + 20.0 * M_PI * 0.01 * joint_vel_filter_freq);
    }

    void updateFilterFreq(T freq)
    {
        joint_vel_filter_freq = freq;
        joint_vel_filter_alpha = 20.0 * M_PI * 0.01 * joint_vel_filter_freq / (1.0 + 20.0 * M_PI * 0.01 * joint_vel_filter_freq);
    }

    void setPD(T kp, T kd)
    {
        joint_kp.setConstant(kp);
        joint_kd.setConstant(kd);
    }

    Eigen::Matrix<T, 3, 3> getRotMat(Eigen::Matrix<T, 4, 1>& quat)
    {
        Eigen::Matrix<T, 3, 3> rot_mat = quatToRotMat(quat);
        return rot_mat;
    }

    Eigen::Matrix<T, 3, 3> getRotMat(Eigen::Matrix<T, 3, 1>& rpy)
    {
        Eigen::Matrix<T, 3, 3> rot_mat = rpyToRotMat<T>(rpy);
        return rot_mat;
    }

    void simpleWcProtect(T tau_max, T v_max, T v_critical)
    {
        // auto tau = (joint_pos_des - joint_pos).cwiseProduct(joint_kp) - joint_vel.cwiseProduct(joint_kd);
        Eigen::Matrix<T, 12, 1> upper, lower;
        upper.setConstant(tau_max);
        lower.setConstant(-tau_max);
        for (int i = 0; i < 12; ++i)
        {
            if (joint_vel[i] > v_critical)
            {
                upper[i] = joint_vel[i] < v_max ? upper[i] / (v_max - v_critical) * (joint_vel[i] - v_critical): 0.0;
            }
            if (joint_vel[i] < -v_critical)
            {
                lower[i] = joint_vel[i] > -v_critical ? lower[i] / (v_max - v_critical) * (- joint_vel[i] - v_critical) : 0.0;
            }
        }
        upper = (upper + joint_vel.cwiseProduct(joint_kd)).cwiseQuotient(joint_kp) + joint_pos;
        lower = (lower + joint_vel.cwiseProduct(joint_kd)).cwiseQuotient(joint_kp) + joint_pos;
        joint_pos_des = joint_pos_des.cwiseMin(upper).cwiseMax(lower);
    }

    void updateContactStatus(int16_t threshold)
    {
        for (int i = 0; i < 4; ++i)
        {
            contact_status.at(i) = foot_force[i] > threshold;
        }
    }

    Eigen::Matrix<T, 12, 1> getFootPosBody()
    {
        auto kin = kinematics<T>();
        Eigen::Matrix<T, 12, 1> output;
        output.setZero();

        T offset_x = 0.1881;
        T offset_y = 0.04675;
        T l_abad = 0.08;
        T l_shank = 0.213;
        T l_thigh = 0.213;

        Eigen::Matrix<T, 3, 1> rho_opt;
        rho_opt.setZero();
        Eigen::Matrix<T, Eigen::Dynamic, 1> rho_fl, rho_fr, rho_hl, rho_hr;
        rho_fl.setZero(5);
        rho_fl << offset_x, offset_y, l_abad, l_shank, l_thigh;
        rho_fr.setZero(5);
        rho_fr << offset_x, -offset_y, -l_abad, l_shank, l_thigh;   
        rho_hl.setZero(5);
        rho_hl << -offset_x, offset_y, l_abad, l_shank, l_thigh;   
        rho_hr.setZero(5);
        rho_hr << -offset_x, -offset_y, -l_abad, l_shank, l_thigh;
        std::array<Eigen::Matrix<T, Eigen::Dynamic, 1>, 4> rho = {rho_fr, rho_fl, rho_hr, rho_hl};

        for (int i = 0; i < 4; ++i)
        {
            output.template segment<3>(3 * i) = kin.fk(joint_pos.template segment<3>(3 * i), rho_opt, rho.at(i));
        }

        return output;
    }

    Eigen::Matrix<T, 12, 1> getFootVelBody()
    {
        auto kin = kinematics<T>();
        Eigen::Matrix<T, 12, 1> output;
        output.setZero();

        T offset_x = 0.1881;
        T offset_y = 0.04675;
        T l_abad = 0.08;
        T l_shank = 0.213;
        T l_thigh = 0.213;

        Eigen::Matrix<T, 3, 1> rho_opt;
        rho_opt.setZero();
        Eigen::Matrix<T, Eigen::Dynamic, 1> rho_fl, rho_fr, rho_hl, rho_hr;
        rho_fl.setZero(5);
        rho_fl << offset_x, offset_y, l_abad, l_shank, l_thigh;
        rho_fr.setZero(5);
        rho_fr << offset_x, -offset_y, -l_abad, l_shank, l_thigh;   
        rho_hl.setZero(5);
        rho_hl << -offset_x, offset_y, l_abad, l_shank, l_thigh;   
        rho_hr.setZero(5);
        rho_hr << -offset_x, -offset_y, -l_abad, l_shank, l_thigh;
        std::array<Eigen::Matrix<T, Eigen::Dynamic, 1>, 4> rho = {rho_fl, rho_fr, rho_hl, rho_hr};

        for (int i = 0; i < 4; ++i)
        {
            auto j_foot = kin.jac(joint_pos.template segment<3>(3 * i), rho_opt, rho.at(i));
            output.template segment<3>(3 * i) = j_foot * joint_vel.template segment<3>(3 * i);
        }
        return output;
    }
};

template <typename T>
class cmdInterface
{
public:
    Eigen::Matrix<T, 3, 1> cmd;
    Eigen::Matrix<T, 3, 1> cmd_range;

    cmdInterface()
    {
        cmd.setZero();
        cmd_range.setZero();
    }

    void setCmdRange(T vx_max, T vy_max, T omega_max)
    {
        cmd_range[0] = vx_max;
        cmd_range[1] = vy_max;
        cmd_range[2] = omega_max;
    }
};